package com.zzw.springboot.util;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;
import java.util.Map;
import java.util.Set;
import java.util.HashSet;
import java.util.HashMap;

/**
 * Redis限流工具类
 */
@Component
public class RedisRateLimiter {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    private static final String RATE_LIMITER_PREFIX = "rate:limiter:";

    /**
     * 判断是否允许访问
     *
     * @param key      限流键（通常使用接口名称或用户ID）
     * @param period   限流周期（秒）
     * @param count 限流周期内最大请求数
     * @return true:允许访问; false:超限,不允许访问
     */
    public boolean isAllowed(String key, int period, int count) {
        key = RATE_LIMITER_PREFIX + key;

        // 查询当前计数器值
        Object countObj = redisTemplate.opsForValue().get(key);
        int currentCount = 0;
        if (countObj != null) {
            currentCount = Integer.parseInt(countObj.toString());
        }

        String ipPassKey = key.replace(RATE_LIMITER_PREFIX, RATE_LIMITER_PREFIX + "pass:");
        String ipLimitedKey = key.replace(RATE_LIMITER_PREFIX, RATE_LIMITER_PREFIX + "limited:");

        // 超过限流阈值
        if (currentCount >= count) {
            // 记录限流请求
            redisTemplate.opsForValue().increment(ipLimitedKey, 1);
            redisTemplate.expire(ipLimitedKey, 24, TimeUnit.HOURS); // 保存24小时用于统计
            return false;
        }

        // 如果是第一次请求，设置过期时间
        if (currentCount == 0) {
            redisTemplate.opsForValue().set(key, 1, period, TimeUnit.SECONDS);
        } else {
            // 计数器加1
            redisTemplate.opsForValue().increment(key, 1);
        }
        // 记录通过请求
        redisTemplate.opsForValue().increment(ipPassKey, 1);
        redisTemplate.expire(ipPassKey, 24, TimeUnit.HOURS); // 保存24小时用于统计
        return true;
    }

    /**
     * 获取所有接口的限流统计数据
     *
     * @return 接口限流统计数据Map
     */
    public Map<String, Map<String, Long>> getInterfaceStats() {
        Map<String, Map<String, Long>> result = new HashMap<>();
        // 获取所有通过和限流的key
        Set<String> passKeys = redisTemplate.keys(RATE_LIMITER_PREFIX + "pass:*");
        Set<String> limitedKeys = redisTemplate.keys(RATE_LIMITER_PREFIX + "limited:*");

        Map<String, Long> passMap = new HashMap<>();
        Map<String, Long> limitedMap = new HashMap<>();

        if (passKeys != null) {
            for (String key : passKeys) {
                String ip = key.substring((RATE_LIMITER_PREFIX + "pass:").length());
                Object count = redisTemplate.opsForValue().get(key);
                if (count != null) {
                    passMap.put(ip, Long.parseLong(count.toString()));
                }
            }
        }
        if (limitedKeys != null) {
            for (String key : limitedKeys) {
                String ip = key.substring((RATE_LIMITER_PREFIX + "limited:").length());
                Object count = redisTemplate.opsForValue().get(key);
                if (count != null) {
                    limitedMap.put(ip, Long.parseLong(count.toString()));
                }
            }
        }
        // 合并所有IP
        Set<String> allIps = new HashSet<>();
        allIps.addAll(passMap.keySet());
        allIps.addAll(limitedMap.keySet());
        for (String ip : allIps) {
            Map<String, Long> stats = new HashMap<>();
            stats.put("passed", passMap.getOrDefault(ip, 0L));
            stats.put("limited", limitedMap.getOrDefault(ip, 0L));
            stats.put("total", stats.get("passed") + stats.get("limited"));
            result.put(ip, stats);
        }
        return result;
    }

}
